import os
import pickle
import json
import numpy as np
from tqdm import tqdm

# 路径设置
pkl_dir = "pymarl/src/offline_data"
json_dir = "offline_data/2s3z/good"
os.makedirs(json_dir, exist_ok=True)


def ensure_list(x):
    """确保转换为标准列表"""
    if hasattr(x, 'tolist'):
        return x.tolist()
    elif isinstance(x, (list, tuple)):
        return list(x)
    else:
        return [float(x)]


for file_name in tqdm(os.listdir(pkl_dir)):
    if not file_name.endswith(".pkl"):
        continue

    file_path = os.path.join(pkl_dir, file_name)
    with open(file_path, "rb") as f:
        episode_data = pickle.load(f)

    serializable_episode = []
    for agent_traj in episode_data:
        traj = []
        for step in agent_traj:
            # 处理所有可能的数据类型
            state = ensure_list(step[0])
            obs = ensure_list(step[1])
            action = [int(step[2][0])] if isinstance(step[2], list) else [int(step[2])]
            reward = [float(step[3][0])] if isinstance(step[3], list) else [float(step[3])]
            done = bool(step[4])
            avail_actions = ensure_list(step[5])

            traj.append([
                state,
                obs,
                action,
                reward,
                done,
                avail_actions
            ])
        serializable_episode.append(traj)

    json_path = os.path.join(json_dir, file_name.replace(".pkl", ".json"))
    with open(json_path, "w") as f:
        json.dump(serializable_episode, f, indent=2)